# -*- coding: utf-8 -*-
import pandas as pd
import talib
import sys
import os

# --- 1. 从命令行获取参数 ---
if len(sys.argv) < 3:
    print("错误：参数不足。")
    print("用法: python3 03_calculate_indicators.py <股票代码> <frequency>")
    sys.exit(1)

stock_code = sys.argv[1]
frequency = sys.argv[2]

# --- 2. 定义输入输出文件名 ---
input_filename = f'output/data/stock_data_{stock_code}_{frequency}.csv'
output_filename = f'output/data/stock_data_{stock_code}_{frequency}_with_indicators.csv'

# --- 3. 加载数据 ---
try:
    print(f"开始加载 {input_filename}...")
    df = pd.read_csv(input_filename)
    print("数据加载完成。")
except FileNotFoundError:
    print(f"错误: 找不到数据文件 {input_filename}。")
    sys.exit(1)

# --- 4. 关键修复：统一使用英文列名 ---
rename_map = {
    '日期': 'date', '股票代码': 'code', '开盘': 'open', '收盘': 'close', 
    '最高': 'high', '最低': 'low', '成交量': 'volume', '成交额': 'amount', 
    '振幅': 'amplitude', '涨跌幅': 'change_percent', '涨跌额': 'change_amount', '换手率': 'turnover'
}
df.rename(columns=rename_map, inplace=True)
print("数据列名已统一为英文。")

# --- 5. 计算技术指标 ---
print("开始计算技术指标...")
close = df['close'].astype(float)
high = df['high'].astype(float)
low = df['low'].astype(float)
volume = df['volume'].astype(float)

df['sma5'] = talib.SMA(close, timeperiod=5)
df['sma10'] = talib.SMA(close, timeperiod=10)
df['sma20'] = talib.SMA(close, timeperiod=20)
df['ema12'] = talib.EMA(close, timeperiod=12)
df['ema26'] = talib.EMA(close, timeperiod=26)
macd, macdsignal, macdhist = talib.MACD(close, fastperiod=12, slowperiod=26, signalperiod=9)
df['macd'] = macd
df['macdsignal'] = macdsignal
df['macdhist'] = macdhist
df['rsi'] = talib.RSI(close, timeperiod=14)
upper, middle, lower = talib.BBANDS(close, timeperiod=20)
df['upperband'] = upper
df['middleband'] = middle
df['lowerband'] = lower
df['atr'] = talib.ATR(high, low, close, timeperiod=14)
df['adx'] = talib.ADX(high, low, close, timeperiod=14)
print("技术指标计算完成。")

# --- 6. 保存数据 ---
# 确保output/data目录存在
os.makedirs('output/data', exist_ok=True)

print(f"开始将结果保存到 {output_filename}...")
df.dropna(inplace=True)
df.to_csv(output_filename, index=False)
print(f"数据保存成功！")
